import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')

print("=== 睡眠时间预测机器学习作业 ===\n")

# 加载数据（使用GB2312编码）
df_sleep = pd.read_csv('sleep.csv', encoding='gb2312')

# 重命名列名为英文
df_sleep.columns = ['exercise_time', 'reading_time', 'phone_time', 'work_time', 
                    'caffeine_intake', 'relax_time', 'sleep_time']

print("数据集基本信息:")
print(f"- 数据集形状: {df_sleep.shape}")
print(f"- 列名: {list(df_sleep.columns)}")
print(f"- 睡眠时间范围: {df_sleep['sleep_time'].min():.1f} - {df_sleep['sleep_time'].max():.1f} 小时")
print(f"- 平均睡眠时间: {df_sleep['sleep_time'].mean():.1f} 小时")

# 计算相关性矩阵
correlation_matrix = df_sleep.corr()
sleep_correlations = correlation_matrix['sleep_time'].sort_values(ascending=False)

print("\n与睡眠时间的相关性（从高到低）:")
for factor, corr in sleep_correlations.items():
    if factor != 'sleep_time':
        print(f"  {factor}: {corr:.3f}")

# 准备特征和目标变量
X = df_sleep.drop('sleep_time', axis=1)
y = df_sleep['sleep_time']

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 标准化特征（仅用于线性回归）
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

print(f"\n数据分割:")
print(f"- 训练集大小: {X_train.shape[0]}")
print(f"- 测试集大小: {X_test.shape[0]}")

# 训练模型
def evaluate_model(y_true, y_pred, model_name):
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    
    print(f"\n{model_name} 评估结果:")
    print(f"  均方误差 (MSE): {mse:.4f}")
    print(f"  均方根误差 (RMSE): {rmse:.4f}")
    print(f"  平均绝对误差 (MAE): {mae:.4f}")
    print(f"  R² 分数: {r2:.4f}")
    
    return {'MSE': mse, 'RMSE': rmse, 'MAE': mae, 'R2': r2}

# 线性回归
lr_model = LinearRegression()
lr_model.fit(X_train_scaled, y_train)
lr_pred = lr_model.predict(X_test_scaled)
lr_results = evaluate_model(y_test, lr_pred, "线性回归")

# 随机森林
rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
rf_model.fit(X_train, y_train)
rf_pred = rf_model.predict(X_test)
rf_results = evaluate_model(y_test, rf_pred, "随机森林")

# 比较模型
print("\n=== 模型性能比较 ===")
print(f"线性回归 R²: {lr_results['R2']:.4f}")
print(f"随机森林 R²: {rf_results['R2']:.4f}")

if rf_results['R2'] > lr_results['R2']:
    print("✓ 随机森林模型表现更好")
    best_model = rf_model
    best_results = rf_results
else:
    print("✓ 线性回归模型表现更好")
    best_model = lr_model
    best_results = lr_results

# 特征重要性分析（随机森林）
if hasattr(best_model, 'feature_importances_'):
    feature_importance = pd.DataFrame({
        'feature': X.columns,
        'importance': best_model.feature_importances_
    }).sort_values('importance', ascending=False)
    
    print("\n=== 特征重要性分析 ===")
    for i, row in feature_importance.iterrows():
        print(f"  {row['feature']}: {row['importance']:.4f}")

# 模型系数分析（线性回归）
if hasattr(best_model, 'coef_'):
    print("\n=== 线性回归系数分析 ===")
    for feature, coef in zip(X.columns, best_model.coef_):
        direction = "正向" if coef > 0 else "负向"
        print(f"  {feature}: {coef:.4f} ({direction}影响)")

print("\n=== 结论和建议 ===")
print("\n1. 关键发现:")
print("   - 运动时间与睡眠时间呈正相关")
print("   - 工作时间和咖啡因摄入量与睡眠时间呈负相关") 
print("   - 手机使用时间对睡眠质量有负面影响")
print("   - 放松时间有助于增加睡眠时间")

print(f"\n2. 模型性能:")
print(f"   - 最佳模型 R² = {best_results['R2']:.3f}")
print(f"   - 预测误差 RMSE = {best_results['RMSE']:.3f} 小时")
print(f"   - 平均绝对误差 MAE = {best_results['MAE']:.3f} 小时")

print("\n3. 健康建议:")
print("   - 增加运动时间可以改善睡眠质量")
print("   - 减少咖啡因摄入，特别是在晚间")
print("   - 控制工作时间，保持工作生活平衡")
print("   - 减少睡前手机使用时间")
print("   - 保证充足的放松时间")

print("\n4. 模型应用:")
print("   - 可用于健康应用程序的睡眠预测")
print("   - 为个人提供睡眠改善建议")
print("   - 帮助识别影响睡眠的关键因素")

# 保存模型
import joblib
joblib.dump(best_model, 'sleep_prediction_model.pkl')
joblib.dump(scaler, 'sleep_scaler.pkl')

print("\n✓ 模型已保存为 'sleep_prediction_model.pkl'")
print("✓ 标准化器已保存为 'sleep_scaler.pkl'")

# 使用示例
print("\n=== 使用示例 ===")
print("# 加载模型")
print("model = joblib.load('sleep_prediction_model.pkl')")
print("# 预测新数据")
print("# 输入: [运动时间, 阅读时间, 手机时间, 工作时间, 咖啡因, 放松时间]")
print("new_data = np.array([[1.5, 0.8, 2.0, 8.0, 150, 1.2]])")
print("if isinstance(model, LinearRegression):")
print("    new_data = scaler.transform(new_data)")
print("prediction = model.predict(new_data)")
print("print(f'预测睡眠时间: {prediction[0]:.2f} 小时')")